from unimport import *
from univeral import *
from auto_binning import *


def calc_psi(x_true, x_pred, num_bins=10, method='freq', lamb=0.001):
    """
    计算特征psi
    :param x_true:
    :param x_pred:
    :param num_bins:
    :param method:
    :return:
    """
    if is_numeric_dtype(x_true):
        bound = base_cut(x_true, num_bins=num_bins, method=method)
        x_base_bins = np.array(make_bin(x_true, bound, num_fillna=-999))
        x_pred_bins = np.array(make_bin(x_pred, bound, num_fillna=-999))
    else:
        x_base_bins = x_true
        x_pred_bins = x_pred

    n_base = len(x_true)
    n_pred = len(x_pred)
    b_bins = np.unique(np.concatenate((x_base_bins, x_pred_bins)))
    psi_mapping = []
    for bin in b_bins:
        c_base = np.count_nonzero(x_base_bins == bin)
        c_pred = np.count_nonzero(x_pred_bins == bin)
        pct_base = (c_base + 0.0) / n_base
        pct_pred = (c_pred + 0.0) / n_pred
        csi = (pct_base - pct_pred) * math.log((pct_base + lamb) / (pct_pred + lamb))
        psi_mapping.append((bin, c_base, c_pred, csi))

    psi = sum([x[3] for x in psi_mapping])
    return psi


def calc_auc(y_true, y_score):
    """
    计算评分或者特征的auc
    :param y_true:
    :param y_score:
    :return:
    """
    return roc_auc_score(y_true, y_score)


def calc_iv(y_true, y_score, num_bins=10, method='freq', fillna=-999, lamb=0.001):
    """
    计算iv
    :param y_true: 标签列
    :param y_score: 特征列
    :return:
    """
    y_true = np.asarray(y_true)
    y_score = np.asarray(y_score)

    if is_numeric_dtype(y_score):
        bound = base_cut(y_score, num_bins=num_bins, method=method)
        y_pred_bins = np.asarray(make_bin(y_score, bound, num_fillna=fillna))
    else:
        bound = chi2merge(y_true, y_score, num_bins=num_bins)
        y_pred_bins = np.asarray(make_bin(y_score, bound))

    N_g = (y_true == 0).sum()
    N_b = (y_true == 1).sum()
    data = np.column_stack((y_true, y_pred_bins))

    iv_arr = []
    bins = np.unique(y_pred_bins)
    for bin in bins:
        data_bin = data[data[:, 1] == bin]
        n_g = (data_bin[:, 0] == 0).sum()
        n_b = (data_bin[:, 0] == 1).sum()
        pct_g = (n_g * 1.0 / N_g)
        pct_b = (n_b * 1.0 / N_b)
        iv = (pct_g - pct_b) * math.log((pct_g + lamb) / (pct_b + lamb))
        iv_arr.append(iv)
    iv = sum(iv_arr)
    return iv


def calc_ks(y_true, y_score):
    """
    计算评分或者特征的ks结果
    :param y_true:
    :param y_score:
    :return:
    """
    y_true = np.asarray(y_true)
    y_score = np.asarray(y_score)
    fpr, tpr, _ = roc_curve(y_true, y_score)
    ks = np.max(np.abs(fpr - tpr))
    return ks


def calc_ks0(y_true, y_score):
    """
    计算评分或者特征的ks
    :param y_true:
    :param y_score:
    :return:
    """
    pred_list = list(y_score)
    label_list = list(y_true)
    total_bad = sum(label_list)
    total_good = len(label_list) - total_bad
    items = sorted(zip(pred_list, label_list), key=lambda x: x[0])
    step = (max(pred_list) - min(pred_list)) / 200
    pred_bin = []
    good_rate = []
    bad_rate = []
    ks_list = []
    for i in range(1, 201):
        idx = min(pred_list) + i * step
        pred_bin.append(idx)
        label_bin = [x[1] for x in items if x[0] < idx]
        bad_num = sum(label_bin)
        good_num = len(label_bin) - bad_num
        goodrate = good_num / total_good
        badrate = bad_num / total_bad
        ks = abs(goodrate - badrate)
        good_rate.append(goodrate)
        bad_rate.append(badrate)
        ks_list.append(ks)
    return max(ks_list)


def fea_psi_calc(actual, predict, bins=10):
    """
    功能: 计算连续变量和离散变量的PSI值
    输入值:
    actual: 一维数组或series，代表训练集中的变量
    predict: 一维数组或series，代表测试集中的变量
    bins: 违约率段划分个数
    输出值:
    字典，键值关系为{'psi': PSI值，'psi_fig': 实际和预期占比分布曲线}
    """
    psi_dict = {}
    actual = np.sort(actual)
    actual_distinct = np.sort(list(set(actual)))
    predict = np.sort(predict)
    # predict_distinct = np.sort(list(set(predict)))
    actual_len = len(actual)
    actual_distinct_len = len(actual_distinct)
    predict_len = len(predict)
    # predict_distinct_len = len(predict_distinct)
    psi_cut = []
    actual_bins = []
    predict_bins = []
    actual_min = actual.min()
    actual_max = actual.max()
    cuts = []
    binlen = (actual_max - actual_min) / bins
    if actual_distinct_len < bins:
        for i in actual_distinct:
            cuts.append(i)
        for i in range(2, (actual_distinct_len + 1)):
            if i == bins:
                lowercut = cuts[i - 2]
                uppercut = float("Inf")
            else:
                lowercut = cuts[i - 2]
                uppercut = cuts[i - 1]
            actual_cnt = ((actual >= lowercut) & (actual < uppercut)).sum() + 1
            predict_cnt = ((predict >= lowercut) & (predict < uppercut)).sum() + 1
            actual_pct = (actual_cnt + 0.0) / actual_len
            predict_pct = (predict_cnt + 0.0) / predict_len
            psi_cut.append(
                (actual_pct - predict_pct) * math.log(actual_pct / predict_pct)
            )
            actual_bins.append(actual_pct)
            predict_bins.append(predict_pct)
    else:
        for i in range(1, bins):
            cuts.append(actual_min + i * binlen)
        for i in range(1, (bins + 1)):
            if i == 1:
                lowercut = float("-Inf")
                uppercut = cuts[i - 1]
            elif i == bins:
                lowercut = cuts[i - 2]
                uppercut = float("Inf")
            else:
                lowercut = cuts[i - 2]
                uppercut = cuts[i - 1]
            actual_cnt = ((actual >= lowercut) & (actual < uppercut)).sum() + 1
            predict_cnt = ((predict >= lowercut) & (predict < uppercut)).sum() + 1
            actual_pct = (actual_cnt + 0.0) / actual_len
            predict_pct = (predict_cnt + 0.0) / predict_len
            psi_cut.append(
                (actual_pct - predict_pct) * math.log(actual_pct / predict_pct)
            )
            actual_bins.append(actual_pct)
            predict_bins.append(predict_pct)
    psi = sum(psi_cut)
    psi_dict["psi"] = psi
    return psi_dict


def cate_univerate(df, feature, target, lamb=0.001, mean_cols=[]):
    feature_code = f"{feature}_code"
    df[feature_code], _ = pd.factorize(df[feature], sort=True)
    t_t = df[target].count()
    b_t = df[target].sum()
    g_t = t_t - b_t
    br = b_t / t_t
    dti = df.groupby([feature, feature_code]).agg(
        bad=(target, 'sum'),
        total=(target, 'count'),
        brate=(target, 'mean')
    )
    dti['total_rate'] = dti['total'] / t_t
    dti['good'] = dti['total'] - dti['bad']
    dti['lift'] = dti['brate'] / br
    for mean_col in mean_cols:
        dti[mean_col] = df.groupby([feature, feature_code])[mean_col].mean()
    dti = dti.dropna().reset_index()

    corr = df[feature_code].corr(df[target], method='spearman')
    # corr, _ = pointbiserialr(df[target], df[feature_code])
    if corr < 0:
        dti = dti.sort_values(feature_code, ascending=False)

    dti["woe"] = np.log(
        ((dti["good"] / g_t) + lamb) / ((dti["bad"] / b_t) + lamb)
    )
    dti['good_cum'] = dti["good"].cumsum()
    dti['bad_cum'] = dti["bad"].cumsum()
    dti['brate_cum'] = dti['bad_cum'] / (dti['good_cum'] + dti['bad_cum'])

    dti["ks"] = np.abs((dti["bad_cum"] / b_t) - (dti["good_cum"] / g_t))
    dti["iv"] = (dti["good"] / g_t - dti["bad"] / b_t) * dti["woe"]
    dti['iv'] = dti['iv'].sum()

    def _cum_calc_auc(n):
        if corr > 0:
            codes_list = range(0, n + 1)
        else:
            max_codes = dti[feature_code].max()
            codes_list = range(n, max_codes + 1)

        df_new = df[df[feature_code].isin(codes_list)]
        auc = metrics.roc_auc_score(df_new[target], df_new[feature_code])
        auc = auc if auc > 0.5 else 1 - auc
        return auc

    dti['auc_cum'] = dti[feature_code].map(_cum_calc_auc)
    dti.insert(0, "target", [target] * dti.shape[0])
    dti.insert(0, "feature", [feature] * dti.shape[0])
    df.drop(columns=[feature_code], inplace=True)
    return dti[['feature', 'target', feature, 'bad', 'total', 'total_rate', 'brate', 'brate_cum', 'lift', 'ks', 'iv',
                'auc_cum'] + mean_cols]


def numerical_univerate(df: pd.DataFrame, feature: str, target: str, labels: list = None, mean_cols: list = [],
                        bins: int = 10, lamb: float = 0.001):
    feature_bin = f"{feature}_bin"
    if labels is None:
        df[feature_bin] = pd.qcut(df[feature], bins, duplicates='drop')
    else:
        df[feature_bin] = pd.cut(df[feature], labels)
    t_t = df[target].count()
    b_t = df[target].sum()
    g_t = t_t - b_t
    br = b_t / t_t
    dti = df.groupby(feature_bin).agg(
        bad=(target, 'sum'),
        total=(target, 'count'),
        brate=(target, 'mean')
    )
    dti['total_rate'] = dti['total'] / t_t
    dti['good'] = dti['total'] - dti['bad']
    dti['lift'] = dti['brate'] / br
    for mean_col in mean_cols:
        dti[mean_col] = df.groupby(feature_bin)[mean_col].mean()
    corr = df[feature].corr(df[target])
    # corr, _ = pointbiserialr(df[target], df[feature])  # 因为是二元分类问题，采用
    if corr < 0:
        dti = dti.sort_values(feature_bin, ascending=False)

    dti["woe"] = np.log(
        ((dti["good"] / g_t) + lamb) / ((dti["bad"] / b_t) + lamb)
    )
    dti['good_cum'] = dti["good"].cumsum()
    dti['bad_cum'] = dti["bad"].cumsum()
    dti['brate_cum'] = dti['bad_cum'] / (dti['good_cum'] + dti['bad_cum'])

    dti["ks"] = np.abs((dti["bad_cum"] / b_t) - (dti["good_cum"] / g_t))
    dti["iv"] = (dti["good"] / g_t - dti["bad"] / b_t) * dti["woe"]
    dti['iv'] = dti['iv'].sum()

    dti = dti.reset_index()
    dti.columns.name = None
    dti['bin_code'] = dti[feature_bin].cat.codes

    def _cum_calc_auc(n):
        if corr > 0:
            codes_list = range(0, n + 1)
        else:
            max_codes = dti['bin_code'].max()
            codes_list = range(n, max_codes + 1)

        df_new = df[df[feature_bin].cat.codes.isin(codes_list)]
        auc = metrics.roc_auc_score(df_new[target], df_new[feature])
        auc = auc if auc > 0.5 else 1 - auc
        return auc

    dti['auc_cum'] = dti['bin_code'].map(_cum_calc_auc)
    dti.rename({feature_bin: 'bin'}, axis=1, inplace=True)
    dti.insert(0, "target", [target] * dti.shape[0])
    dti.insert(0, "feature", [feature] * dti.shape[0])
    return dti[['feature', 'target', 'bin', 'bad', 'total', 'total_rate', 'brate', 'brate_cum', 'lift', 'ks', 'iv',
                'auc_cum'] + mean_cols]
